#!/usr/bin/env python3

"""
Module for querying a SpecDB database
"""

import csv
import io
import json
import os
import re
import sys
import sqlite3
import tarfile
import time

import pandas as pd
from specdb.STAR import fid2star


def write_out_directory(dic=None, head_dir=None):
	
	date_info = dic['experiment_date'].split()
	date      = date_info[0].split('-')
	date      = ''.join(date)
	
	ids = [dic['pst_preparer'], dic['pst_id'], date]
	
	dir_name = '_'.join(ids)
	
	if not os.path.isdir(os.path.join(head_dir, dir_name)):
		os.mkdir(os.path.join(head_dir, dir_name))
	
	f = io.BytesIO()
	tar_file = tarfile.open
	
	f = io.BytesIO()
	tar_file = tarfile.open(fileobj=f, mode='w:')
	tar_info = tarfile.TarInfo('data_zipped')
	tar_info.size = len(dic['zipped_dir'])
	tar_info.mtime = time.time()
	tar_file.addfile(tar_info, io.BytesIO(dic['zipped_dir']))
	tar_file.close()
	os.chdir(os.path.join(head_dir, dir_name))
	
	with open('stream.tar.gz', 'wb') as fp:
		fp.write(f.getvalue())
	
	os.system("tar -xf stream.tar.gz")
	os.system("tar -xf data_zipped")
	
	os.remove("stream.tar.gz")
	os.remove("data_zipped")
	
	return dir_name


def query_from_sql(cursor=None, sql=None, star=None, output_dir=None):
	"""
	Query from user provided SQL SELECT statement
	"""
	
	# execute query
	cursor.execute(sql)
	
	results = cursor.fetchall()
	results = [dict(res) for res in results]
	
	if output_dir is None:
		df = pd.DataFrame(results)
		pd.set_option('display.max_columns', None)
		pd.set_option('max_colwidth', 20)
		pd.set_option('display.max_rows', None)
		pd.set_option('display.width', 110)
		pd.set_option('max_rows', None)
		
		blankindex = [''] * len(df)
		df.index = blankindex
		df = df.rename(columns=lambda x: x[:14])
		
		print()
		print(df)
		print()
		
		return
	
	if not os.path.isdir(os.path.abspath(output_dir)):
		os.mkdir(os.path.abspath(output_dir))
	
	output_dir = os.path.abspath(output_dir)
	
	for res in results:
		
		id_dir = write_out_directory(dic=res, head_dir=output_dir)
		
		sql =  f"SELECT subdir_name FROM time_domain_dataset WHERE"
		sql += f" id='{res['id']}'"
		
		cursor.execute(sql)
		subname = cursor.fetchall()[0]
		#print(subname['subdir_name'])
		if star:
			star_path = os.path.join(
				output_dir,
				id_dir,
				subname['subdir_name'],
				'fid.star'
			)
			#print(star_path)
			out = fid2star(
				data=res,
				cursor=cursor,
				save=star_path
			)
	
	return True


def query_from_ids(cursor=None, ids=None, star=None, output_dir=None):
	"""
	Select rows from list of ids
	"""
	
	all_results = []
	results = []
	
	sql = f"SELECT * FROM summary WHERE id='{ids[0]}'"
	#print(ids)
	for i in range(1,len(ids)):
		sql += f" OR id='{ids[i]}'"
	
	#print(sql)
	query_from_sql(
		cursor=cursor,
		sql=sql,
		star=star,
		output_dir=output_dir
	)
	
	return True
	
	
def query_from_file(cursor=None, ids_file=None, star=None, output_dir=None):
	"""
	Collect ids from a file to query against database
	"""
	
	with open(ids_file, 'r') as fp:
		
		ids = dict()
		reader = csv.reader(fp)
		
		csv_data = list(reader)
		
		for row in csv_data:
			for id in row:
				if id not in ids: ids[id] = True
		
		ids = list(ids.keys())
	
	query_from_ids(cursor=cursor, ids=ids, star=star, output_dir=output_dir)
	return True
	

def query(db=None, sql=False, indices=False, star=False, output_dir=None):
	"""
	Perform a query against summary view of database
	Format the query results either in a directory hierarchy or in NMR-STAR
	for each retrieved FID
	
	Parameters
	----------
	+ db			path to a specdb database file
	+ sql			user provided SQL select statement, must be on summary
	+ star			whether to output a STAR file for each FID in the query
	+ indices		list/file of row ids to select from SpecDB database
	+ output_dir	where to put the results
	
	Returns
	-------
	Directory structure of results with NMR-STAR at the end for each FID
	"""
	
	if db is None:
		print("Must provide a database file")
		print("Aborting")
		sys.exit()
	
	if not sql and not indices:
		print("Must provide a SQL SELECT statement or a list of ids")
		print("Cannot provide neither")
		print("Aborting")
		sys.exit()
	
# 	if output_dir is None:
# 		print("Must provide an output directory to place query results")
# 		print("Aborting")
# 		sys.exit()
	
	# connect to database
	conn = sqlite3.connect(db)
	conn.row_factory = sqlite3.Row
	c = conn.cursor()
	#print(star)
	if not indices:
		if type(sql) != str:
			print(f'Unexpected type for argument `sql` type: {type(sql)}')
			sys.exit()
		
		if not re.search(r'^SELECT', sql):
			print(f"provided SQL does not look like a SELECT statement")
			print(f"provided SQL: {sql}")
			print("Aborting")
			sys.exit()
	
		if not re.search(r' summary ', sql):
			print("provided SQL does not appear to SELECT on summary")
			print(f" provided SQL: {sql}")
			print("Aborting")
			sys.exit()
		
		query_from_sql(
			cursor=c,
			sql=sql,
			star=star,
			output_dir=output_dir
		)
	
	if not sql:
		if type(indices) == str and indices.endswith('.csv'):
			query_from_file(
				cursor=c,
				ids_file=indices,
				star=star,
				output_dir=output_dir
			)
		else:
			#print(indices)
			query_from_ids(
				cursor=c,
				ids=indices,
				star=star,
				output_dir=os.path.abspath(output_dir)
			)

